import cv2
import threading
import time
import os
import tkinter as tk
from tkinter import ttk, messagebox, filedialog
import datetime
import logging
import queue
from PIL import Image, ImageTk  # 明确导入PIL模块

class ESP32CAMVideoRecorder:
    def __init__(self, root):
        self.root = root
        self.root.title("ESP32-CAM视频流录制器")
        self.root.geometry("800x600")
        self.root.resizable(True, True)

        # 设置中文字体
        self.style = ttk.Style()
        self.style.configure("TLabel", font=("SimHei", 10))
        self.style.configure("TButton", font=("SimHei", 10))
        self.style.configure("TEntry", font=("SimHei", 10))

        # 创建日志目录
        if not os.path.exists("logs"):
            os.makedirs("logs")

        # 配置日志
        self.logger = logging.getLogger("ESP32CAMRecorder")
        self.logger.setLevel(logging.INFO)
        formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")

        # 文件日志
        file_handler = logging.FileHandler(
            os.path.join(
                "logs",
                f"recorder_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.log",
            )
        )
        file_handler.setFormatter(formatter)
        self.logger.addHandler(file_handler)

        # 控制台日志
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(formatter)
        self.logger.addHandler(console_handler)

        # 视频流参数
        self.stream_url = tk.StringVar(value="http://192.168.100.103/mjpeg/1")
        self.output_dir = tk.StringVar(value=os.path.join(os.getcwd(), "videos"))
        self.frame_queue = queue.Queue(maxsize=100)  # 队列大小可根据需求调整
        self.is_recording = False
        self.is_streaming = False
        self.stream_thread = None
        self.record_thread = None
        self.last_frame = None  # 保存最后一帧用于填充
        self.target_fps = 24  # 固定目标帧率
        self.frame_interval = 1.0 / self.target_fps  # 每帧时间间隔（约0.0417秒）
        self.start_time = 0.0  # 录制开始时间
        self.frame_count = 0  # 已写入帧数
        self.current_file = ""

        # 创建GUI组件
        self.create_widgets()

        self.logger.info("程序已启动")

    def create_widgets(self):
        # 连接设置框架
        conn_frame = ttk.LabelFrame(self.root, text="连接设置")
        conn_frame.pack(fill="x", padx=10, pady=5)

        ttk.Label(conn_frame, text="视频流URL:").grid(
            row=0, column=0, padx=5, pady=5, sticky="w"
        )
        ttk.Entry(conn_frame, textvariable=self.stream_url, width=50).grid(
            row=0, column=1, padx=5, pady=5
        )

        ttk.Label(conn_frame, text="保存目录:").grid(
            row=1, column=0, padx=5, pady=5, sticky="w"
        )
        save_dir_frame = ttk.Frame(conn_frame)
        save_dir_frame.grid(row=1, column=1, padx=5, pady=5, sticky="ew")

        ttk.Entry(save_dir_frame, textvariable=self.output_dir, width=42).pack(
            side="left", fill="x", expand=True
        )
        ttk.Button(save_dir_frame, text="浏览", command=self.browse_output_dir).pack(
            side="left", padx=5
        )

        # 控制按钮框架
        control_frame = ttk.Frame(self.root)
        control_frame.pack(fill="x", padx=10, pady=5)

        self.start_btn = ttk.Button(
            control_frame, text="开始录制", command=self.start_recording
        )
        self.start_btn.pack(side="left", padx=5)

        self.stop_btn = ttk.Button(
            control_frame,
            text="停止录制",
            command=self.stop_recording,
            state="disabled",
        )
        self.stop_btn.pack(side="left", padx=5)

        self.status_label = ttk.Label(
            control_frame, text="状态: 未录制", foreground="red"
        )
        self.status_label.pack(side="left", padx=20)

        self.fps_label = ttk.Label(control_frame, text=f"FPS: {self.target_fps}")
        self.fps_label.pack(side="left", padx=20)

        # 视频预览框架
        preview_frame = ttk.LabelFrame(self.root, text="视频预览")
        preview_frame.pack(fill="both", expand=True, padx=10, pady=5)

        self.preview_label = ttk.Label(preview_frame, text="等待视频流...")
        self.preview_label.pack(fill="both", expand=True, padx=10, pady=10)

        # 日志框架
        log_frame = ttk.LabelFrame(self.root, text="日志")
        log_frame.pack(fill="x", padx=10, pady=5)

        self.log_text = tk.Text(log_frame, height=5, state="disabled")
        self.log_text.pack(fill="x", padx=5, pady=5)

        scrollbar = ttk.Scrollbar(self.log_text, command=self.log_text.yview)
        scrollbar.pack(side="right", fill="y")
        self.log_text.config(yscrollcommand=scrollbar.set)

    def browse_output_dir(self):
        directory = filedialog.askdirectory(title="选择保存目录")
        if directory:
            self.output_dir.set(directory)
            if not os.path.exists(directory):
                os.makedirs(directory)
                self.log_message(f"创建目录: {directory}")

    def log_message(self, message):
        self.logger.info(message)
        self.log_text.config(state="normal")
        self.log_text.insert("end", f"{time.strftime('%H:%M:%S')} - {message}\n")
        self.log_text.see("end")
        self.log_text.config(state="disabled")

    def start_recording(self):
        if self.is_recording:
            return

        # 确保保存目录存在
        save_dir = self.output_dir.get()
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
            self.log_message(f"创建目录: {save_dir}")

        self.is_recording = True
        self.is_streaming = True
        self.frame_count = 0
        self.start_time = time.time()  # 记录开始时间
        self.last_frame = None  # 重置最后一帧

        # 更新UI状态
        self.start_btn.config(state="disabled")
        self.stop_btn.config(state="normal")
        self.status_label.config(text="状态: 正在录制", foreground="green")
        self.fps_label.config(text=f"FPS: {self.target_fps}")

        # 启动流处理线程
        self.stream_thread = threading.Thread(target=self.process_stream)
        self.stream_thread.daemon = True
        self.stream_thread.start()

        # 启动录制线程
        self.record_thread = threading.Thread(target=self.record_frames)
        self.record_thread.daemon = True
        self.record_thread.start()

        self.log_message(f"开始从 {self.stream_url.get()} 录制视频（FPS={self.target_fps}）")

    def stop_recording(self):
        if not self.is_recording:
            return

        self.is_recording = False
        self.is_streaming = False

        # 清空队列
        while not self.frame_queue.empty():
            try:
                self.frame_queue.get_nowait()
            except queue.Empty:
                continue

        # 更新UI状态
        self.start_btn.config(state="normal")
        self.stop_btn.config(state="disabled")
        self.status_label.config(text="状态: 未录制", foreground="red")

        self.log_message("停止录制")

    def process_stream(self):
        try:
            # 打开视频流
            cap = cv2.VideoCapture(self.stream_url.get())

            if not cap.isOpened():
                self.log_message("无法打开视频流")
                self.is_recording = False
                self.is_streaming = False
                self.root.after(1000, self.update_ui_after_error)
                return

            self.log_message("成功连接到视频流")

            while self.is_streaming:
                ret, frame = cap.read()

                if not ret:
                    self.log_message("读取帧失败，可能视频流已断开")
                    self.is_streaming = False
                    break

                # 调整帧大小以适应预览窗口
                height, width = frame.shape[:2]
                max_width = 640
                max_height = 480

                if width > max_width or height > max_height:
                    ratio = min(max_width / width, max_height / height)
                    frame = cv2.resize(frame, (int(width * ratio), int(height * ratio)))

                # 队列满时丢弃最旧帧，保留最新帧
                try:
                    self.frame_queue.put_nowait(frame)
                except queue.Full:
                    try:
                        self.frame_queue.get_nowait()
                        self.frame_queue.put_nowait(frame)
                    except queue.Empty:
                        pass

                self.last_frame = frame  # 始终保存最新帧
                self.update_preview(frame)

                # 检查流是否中断（超过2秒无新帧）
                if self.last_frame is not None and (time.time() - self.start_time > 2 and self.frame_count == 0):
                    self.log_message("视频流初始连接延迟，正在等待帧数据...")

                # 短暂休眠，避免CPU占用过高
                time.sleep(0.001)  # 轻微休眠减少空转

            cap.release()
            self.log_message("视频流已关闭")

        except Exception as e:
            self.log_message(f"处理视频流时发生错误: {str(e)}")
            self.is_streaming = False
            if self.is_recording:
                self.root.after(1000, self.update_ui_after_error)

    def update_preview(self, frame):
        # 将OpenCV帧转换为Tkinter可用的格式
        try:
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            image = Image.fromarray(rgb_frame)
            photo = ImageTk.PhotoImage(image=image)
            self.preview_label.config(image=photo)
            self.preview_label.image = photo
        except Exception as e:
            self.log_message(f"更新预览时发生错误: {str(e)}")

    def update_ui_after_error(self):
        self.is_recording = False
        self.start_btn.config(state="normal")
        self.stop_btn.config(state="disabled")
        self.status_label.config(text="状态: 未录制", foreground="red")

    def record_frames(self):
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        out = None
        frame_size = None

        try:
            while self.is_recording:
                # 计算已录制时间和应写入的帧数
                elapsed_time = time.time() - self.start_time
                expected_frames = int(elapsed_time * self.target_fps) + 1  # 向上取整保证帧数足够

                # 按固定帧率循环写入
                while self.frame_count < expected_frames:
                    if self.is_recording:
                        # 检查队列是否有新帧
                        if not self.frame_queue.empty():
                            frame = self.frame_queue.get()
                            self.last_frame = frame  # 更新最新帧
                        else:
                            # 队列为空，使用最后一帧填充
                            if self.last_frame is None:
                                # 初始化前无帧，跳过（避免写入空帧）
                                time.sleep(self.frame_interval)
                                continue
                            frame = self.last_frame

                        # 初始化视频写入器
                        if out is None:
                            height, width = frame.shape[:2]
                            frame_size = (width, height)
                            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                            output_file = os.path.join(self.output_dir.get(), f"recording_{timestamp}.mp4")
                            self.current_file = output_file
                            out = cv2.VideoWriter(output_file, fourcc, self.target_fps, frame_size)
                            self.log_message(f"开始写入视频文件: {output_file} (FPS={self.target_fps})")

                        # 写入帧
                        out.write(frame)
                        self.frame_count += 1

                    # 控制帧率间隔
                    sleep_time = max(0, (self.frame_count / self.target_fps) - elapsed_time)
                    time.sleep(sleep_time)

        except Exception as e:
            self.log_message(f"录制过程中发生错误: {str(e)}")
        finally:
            # 确保资源释放
            if out is not None:
                out.release()
                if os.path.exists(self.current_file):
                    file_size = os.path.getsize(self.current_file) / (1024 * 1024)
                    self.log_message(f"视频文件已保存: {self.current_file}, 大小: {file_size:.2f} MB")
                else:
                    self.log_message(f"视频文件 {self.current_file} 未正确保存。")

            # 更新UI状态（流中断处理）
            if not self.is_streaming and self.is_recording:
                self.root.after(
                    0,
                    lambda: self.status_label.config(
                        text="状态: 流已断开，文件已保存", foreground="orange"
                    ),
                )
                self.root.after(
                    3000,
                    lambda: self.status_label.config(
                        text="状态: 未录制", foreground="red"
                    ),
                )
                self.root.after(0, lambda: self.start_btn.config(state="normal"))
                self.root.after(0, lambda: self.stop_btn.config(state="disabled"))
                self.is_recording = False


if __name__ == "__main__":
    root = tk.Tk()
    app = ESP32CAMVideoRecorder(root)
    root.mainloop()